import cv2
import argparse
import os
import numpy as np
import sys
from pathlib import Path
 
sys.path.append(str(Path(__file__).resolve().parent.parent))

from file_util.file_util import GetAllSubFile
import xml.etree.ElementTree as ET

def draw_image_with_anno(new_img_file,new_anno_file):
		try:
			tree = ET.parse(new_anno_file)
			root = tree.getroot() 
		except:
			print(new_anno_file + " is not exist.")
			return
		img = cv2.imread(new_img_file)
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				cv2.rectangle(img,(xmin,ymin),(xmax,ymax),(0,0,255),5)
		cv2.imwrite(new_img_file,img)


def update_anno(anno_file, new_anno_file):
		try:
			tree = ET.parse(anno_file)
			root = tree.getroot() 
		except:
			print(anno_file + " is not exist.")
			return
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)

				new_xmin = xmin
				new_ymin = ymin
				new_xmax = xmax
				new_ymax = ymax
				object[4][0].text = str(new_xmin)
				object[4][1].text = str(new_ymin)
				object[4][2].text = str(new_xmax)
				object[4][3].text = str(new_ymax)
		tree.write(new_anno_file,encoding="utf-8", xml_declaration=True)
		return 

def img_blur(opt,kernel_size_list):
	img_files = GetAllSubFile(opt.root_data)
	for img_file in img_files:
		img = cv2.imread(img_file)
		h = img.shape[0]
		w = img.shape[1]
		path, img_file_name = os.path.split(img_file)
		img_name = img_file_name.rsplit('.',1)[0]
		#draw_totation(img_file,os.path.join(opt.anon_dir,img_name + '.xml'))
		anno_file = os.path.join(opt.anon_dir,img_name + '.xml')
		for ksize in kernel_size_list:
			blur_img = cv2.blur(img, (ksize,ksize))
			new_img_name = img_name+"_blur_" + str(ksize)
			new_anno_file = os.path.join(opt.anon_dir,new_img_name + '.xml')
			new_img_file = os.path.join(path,new_img_name +".jpg")
			cv2.imwrite(new_img_file,blur_img)
			print(new_anno_file)
			update_anno(anno_file, new_anno_file)
			#draw_image_with_anno(new_img_file,new_anno_file)

if __name__ == '__main__':
   parser = argparse.ArgumentParser()
   parser.add_argument('--root_data', type=str, default='./data', help='root data dir.')
   parser.add_argument('--anon_dir', type=str, default='./anno', help='anno dir.')
   opt = parser.parse_args()
   opt.anon_dir = opt.anon_dir.strip()
   kernel_size_list = [7,13,17]
   #3
   img_blur(opt,kernel_size_list)